from utils import print_timing

def decode_genotype(call, alleleA, alleleB):
    """Returns genotype value corresponding to genotype call (number format) and allele."""
    decode_dict = {'-1':'\t--',
                   '0':'\t'+alleleA+alleleA,
                   '1':'\t'+alleleA+alleleB,
                   '2':'\t'+alleleB+alleleB}
    return decode_dict[call]

def hwe(p,q):
    q = 1-q
    out = [(1-p)**2*q,2*p*(1-p)*q,p**2*q,0]
    out[3] = 1-sum(out)
    return out

# This function is copied from anno.py
def discard_lines_before_annotate_header(annotfile):
    import string, sys
    af = open(annotfile,'r')
    discard_lines = 0
    discard_chars = string.whitespace + ','
    while True:
        line = af.readline()
        line_split = line.lstrip(discard_chars).rstrip(discard_chars).split(',')
        if(len(line_split) > 10):
            break
        if not line or discard_lines > 10000:
            sys.exit("""Couldn't find annotate header in the first %s lines of %s.
                        Is it a correctly formatted annotation file?
                     """%(discard_lines, annotfilename))
        discard_lines += 1
    af.close()
    return discard_lines

def getpids(platform, annotfile, K, truncannotfile=None):
    import os
    from itertools import takewhile
    from utils import file_exists

    # if K is None:
    #     affy6cond = lambda x: True
    # else:
    #     affy6cond = lambda x: x[0] < K

    k_is_none = K is None
    if K is None:
        illuminacond = lambda x: x.split(',')[0] != '[Controls]'
        def action(line): pass
    else:
        illuminacond = lambda x: x.split(',')[0] != '[Controls]' and i < K
        def action(line): tf.write(line)
    if file_exists(annotfile, False):
        af = open(annotfile,'r')
    if K is None and truncannotfile is not None:
        print "Since K is not specified, the original annotation file %s\nwill be the same as the truncated version\n"%(annotfile) \
              + "So not generating truncated annotation file %s"%(truncannotfile)
    if K is not None:
        if truncannotfile is None:
            annotfile_basename = os.path.basename(annotfile)
            truncannotfile = annotfile_basename.rpartition('.')[0] + '-K' + str(K) + '.' + annotfile_basename.rpartition('.')[2]
            print "writing truncated annotation file: name not specified, using %s"%truncannotfile
        else:
            print "writing truncated annotation file %s"%truncannotfile
        tf=open(truncannotfile, 'w')

    for i in range(discard_lines_before_annotate_header(annotfile) + 1):
        line = af.readline()
        if K is not None:
            tf.write(line)
    pids = []
    if platform=='affy6':
        if K is None:
            for line in af:
                pids.append(line.split(',')[0].strip('"')+'\n')
        else:
            for i, line in takewhile(lambda x: x[0] < K, enumerate(af)):
                pids.append(line.split(',')[0].strip('"')+'\n')
                tf.write(line)
    elif platform == 'illumina':
        i = 0
        for line in takewhile(illuminacond, af):
            line_split=line.split(',')
            pid=line_split[1][0:3]
            out = line_split[1] + ',' + line_split[2] + ',' + line_split[3][1] + line_split[3][3] + ',' \
                  + line_split[15] + ',' + line_split[9] + ',' + line_split[10]
            if pid!='cnv' and pid!='hCV' and pid!='cnv':
                i = i+1
                pids.append(out.strip('"')+'\n')
                action(line)
    af.close()
    if K is not None:
        tf.close()
    #of.close()
    return pids

def find_annotfile(platform):
    from utils import get_conf
    conf = get_conf()
    from utils import file_exists
    if "annotfile" in conf[platform] and file_exists(conf[platform]["annotfile"], False):
        annotfile = conf[platform]["annotfile"]
        print "attempting to use annotation file %s specified in config"%annotfile
    else:
        if platform == 'affy6':
            annotfile = "data/affy-anno.txt"
            print "using file '%s' as annotation file"%annotfile
        elif platform == 'illumina':
            annotfile = "data/illumina-anno.txt"
            print "using file '%s' as annotation file"%annotfile
        else:
            sys.exit("sorry, can't find any usable annotation files")
    return annotfile

@print_timing
def simdat(platform, samplesize, annotfile, destdir, K, oname, seedval, truncannotfile=None):
    import sys
    if platform != 'affy6' and platform != 'illumina':
        sys.exit("platform must be either 'affy6' or 'illumina', not %s"%platform)
    from numpy.random import beta, binomial, exponential, multinomial, normal, random_integers, uniform
    from random import sample
    from numpy.random import seed as npseed
    from random import seed as randseed
    npseed(seedval)
    randseed(seedval)
    import string, glob, os
    import utils

    import os, sys
    #from utils import file_exists
    if platform != 'affy6' and platform != 'illumina':
        sys.exit("platform must be either 'affy6' or 'illumina', not %s"%platform)

    # attempt to use annotfile from config if none is specified
    if annotfile == None:
        annotfile = find_annotfile(platform)
    annotfile = os.path.abspath(annotfile)

    revcomp = {'C':'G','G':'C','A':'T','T':'A'}
    anno = []
    pids = getpids(platform, annotfile, K, truncannotfile)
    for pid in pids:
       anno.append(pid.strip('\n'))
    if K == None:
        annonum = len(anno)
    else:
        annonum = min(K, len(anno))
    gtfname = os.path.join(destdir, oname + '-n' + str(samplesize) + '-K' + str(annonum) + '-' + platform + '-geno.txt')
    phfname = os.path.join(destdir, oname + '-n' + str(samplesize) + '-K' + str(annonum) + '-' + platform + '-pheno.txt')
    gtfile = open(gtfname,'w')
    phfile = open(phfname,'w')
    phfile.write('expid, famid, patid, pid, mid, sex, pheno, bp, ostime, osevent, race\n')
    raceconv = {0:'White', 1:'Black', 2:'Asian', 3:'Unknown'}
    gender = {0:'1', 1:'2'}

    print "writing pheno file %s"%phfname

    for i in range(samplesize):
        ii = str(i+1)
        patid = 'PATID' + ii
        sampid = 'SAMPID' + ii
        expid = 'SAMPLE-' + str(ii)
        if platform == 'affy6':
            expid = expid+'.CEL'
        sex = gender[random_integers(0, 1)]
        pheno = str(binomial(1, 0.5))
        bp = str(normal(120, 10))
        ostime = str(exponential(1))
        osevent = str(binomial(1, 0.8))
        race = raceconv[random_integers(0, 3)]
        famid = "FAMID" + ii
        pid = "0"
        mid = "0"
        phfile.write(expid + ',' + famid + ',' + patid + ',' + pid + ',' + mid + ',' + sex + ',' + pheno + ','+bp + ',' + ostime + ',' + osevent + ',' + race+'\n')
    phfile.close()

    print "writing geno file %s"%gtfname

    if platform == 'illumina':
             gtfile.write('[Data]\n')
             gtfile.writelines('' + '\t' + '\t'.join(["SAMPLE-" + str(j) for j in range(1, samplesize+1)]) + '\n')
    elif platform == 'affy6':
        gtfile.writelines('probeset_id' + '\t' + '\t'.join(["SAMPLE-"+str(j) + ".CEL" for j in range(1, samplesize+1)])+'\n')
    for i in range(annonum):
        annodat = anno[i].split(',')
        pid = annodat[0]
        gtypes = multinomial(samplesize, hwe(uniform(),beta(1, 2)))
        gtypes = gtypes[0]*['0']+gtypes[1]*['1']+gtypes[2]*['2']+gtypes[3]*['-1']
        gtypes = sample(gtypes, samplesize)
        if platform == 'illumina':
            istrand = annodat[1]
            A = annodat[2][0]
            B = annodat[2][1]
            sstrand = annodat[3]
            if(istrand != sstrand):
                A = revcomp[A]
                B = revcomp[B]
            gtfile.write(pid)
            for i in range(samplesize):
                gtfile.write(decode_genotype(gtypes[i], A, B))
        if platform == 'affy6':
            x = [pid] + gtypes
            gtfile.writelines('\t'.join(x))
        if i < annonum:
            gtfile.write('\n')
    gtfile.close()

if __name__ == '__main__':
    import sys
    from optparse import OptionParser
    usage = "usage: %prog [-d destdir -k K -o oname] platform samplesize"
    parser = OptionParser(usage=usage)
    parser.add_option("-a", "--annotfile", action="store", type="string", dest="annotfile", help="annotation file name")
    parser.add_option("-d","--destdir", action="store", dest="destdir", type="string", default=".", help="destination directory. Default is current directory.")
    parser.add_option("-k", action="store", dest="K", type="int", help="maximum number of probes to be processed")
    parser.add_option("-o", "--outname", action="store", dest="oname", default="dummy", help="prefix of output files (idlink, pheno and geno)")
    parser.add_option("-t", action="store", dest="truncannotfile", type="string", help="truncate annotation file")
    (options, args) = parser.parse_args()

    if len(args) != 2:
        print "wrong number of arguments"
        parser.print_help()
        sys.exit(1)

    platform = args[0]
    samplesize = int(args[1])
    simdat(platform, samplesize, options.annotfile, options.destdir, options.K, options.oname, options.truncannotfile)
